#include "aes.h"
#include <openssl/modes.h>
#include <openssl/aes.h>

int getPKCS7PaddedLength(int dataLen, int alignSize)
{
    // 计算填充的字节数（按alignSize字节对齐进行填充）
    int remainder = dataLen % alignSize;
    int paddingSize = (remainder == 0) ? alignSize : (alignSize - remainder);
    return (dataLen + paddingSize);
}

/**
 * @brief PKCS7Padding
 * 采用PKCS7Padding方式，将in数据进行alignSize字节对齐填充。
 * 此函数用于加密前，对明文进行填充。
 * @param in 数据
 * @param alignSize 对齐字节数
 * @return 返回填充后的数据
 */
std::string PKCS7Padding(const std::string &in, int alignSize)
{
    // 计算需要填充字节数（按alignSize字节对齐进行填充）
    int remainder = in.size() % alignSize;
    int paddingSize = (remainder == 0) ? alignSize : (alignSize - remainder);

    // 进行填充
    std::string temp(in);
    temp.append(paddingSize, paddingSize);
    return temp;
}

/**
 * @brief PKCS7UnPadding
 * 采用PKCS7Padding方式，将in数据去除填充。
 * 此函数用于解密后，对解密结果进一步去除填充，以得到原始数据。
 * @param in 数据
 * @return 返回去除填充后的数据
 */
std::string PKCS7UnPadding(const std::string &in)
{
    char paddingSize = in.at(in.size() - 1);
    return in.substr(0, in.size() - paddingSize);
}


std::string DoPad(const std::string &in, bool enc, Aes::Padding padding)
{
    if (enc)
    {
        switch (padding)
        {
        case Aes::PADDING_PKCS7:
            // 进行PKCS7Padding填充
            return PKCS7Padding(in, AES_BLOCK_SIZE);
        default:
            return in;
        }
    }
    else
    {
        switch (padding)
        {
        case Aes::PADDING_PKCS7:
            // 解除PKCS7Padding填充
            return PKCS7UnPadding(in);
        default:
            return in;
        }
    }
}


bool Aes::ecb_encrypt(const std::string &in, std::string &out, const std::string &key, bool enc, Padding padding)
{    
    // 检查密钥合法性(只能是16、24、32字节)
    if (key.size() != 16 && key.size() != 24 && key.size() != 32)
    {
        return false;
    }

    if (enc)
    {
        // 生成加密key
        AES_KEY aes_key;
        if (AES_set_encrypt_key((const unsigned char*)key.data(), key.size() * 8, &aes_key) != 0)
        {
            return false;
        }

        std::string inTemp = DoPad(in, enc, padding);

        // 执行ECB模式加密
        out.resize(inTemp.size()); // 调整输出buf大小
        for (int i = 0; i < inTemp.size() / AES_BLOCK_SIZE; i++)
        {
            AES_ecb_encrypt((const unsigned char*)inTemp.data() + AES_BLOCK_SIZE * i,
                            (unsigned char*)out.data() + AES_BLOCK_SIZE * i,
                            &aes_key,
                            AES_ENCRYPT);
        }
        return true;
    }
    else
    {
        // 生成解密key
        AES_KEY aes_key;
        if (AES_set_decrypt_key((const unsigned char*)key.data(), key.size() * 8, &aes_key) != 0)
        {
            return false;
        }

        // 执行ECB模式解密
        out.resize(in.size()); // 调整输出buf大小
        for (int i = 0; i < in.size() / AES_BLOCK_SIZE; i++)
        {
            AES_ecb_encrypt((const unsigned char*)in.data() + AES_BLOCK_SIZE * i,
                            (unsigned char*)out.data() + AES_BLOCK_SIZE * i,
                            &aes_key,
                            AES_DECRYPT);
        }

        out = DoPad(out, enc, padding);
        return true;
    }
}

bool Aes::cbc_encrypt(const std::string &in, std::string &out, const std::string &key, const std::string &ivec, bool enc, Padding padding)
{
    // 检查密钥合法性(只能是16、24、32字节)
    if (key.size() != 16 && key.size() != 24 && key.size() != 32)
    {
        return false;
    }
    if (ivec.size() != 16) // 初始向量为16字节
    {
        return false;
    }

    if (enc)
    {
        // 生成加密key
        AES_KEY aes_key;
        if (AES_set_encrypt_key((const unsigned char*)key.data(), key.size() * 8, &aes_key) != 0)
        {
            return false;
        }

        std::string inTemp = DoPad(in, enc, padding);

        // 执行CBC模式加密
        std::string ivecTemp(ivec.data(), ivec.size()); // ivec会被修改，故需要临时变量来暂存
        out.resize(inTemp.size()); // 调整输出buf大小
        AES_cbc_encrypt((const unsigned char*)inTemp.data(),
                        (unsigned char*)out.data(),
                        inTemp.size(),
                        &aes_key,
                        (unsigned char*)ivecTemp.data(),
                        AES_ENCRYPT);
        return true;
    }
    else
    {
        // 生成解密key
        AES_KEY aes_key;
        if (AES_set_decrypt_key((const unsigned char*)key.data(), key.size() * 8, &aes_key) != 0)
        {
            return false;
        }

        // 执行CBC模式解密
        std::string ivecTemp(ivec.data(), ivec.size()); // ivec会被修改，故需要临时变量来暂存
        out.resize(in.size()); // 调整输出buf大小
        AES_cbc_encrypt((const unsigned char*)in.data(),
                        (unsigned char*)out.data(),
                        in.size(),
                        &aes_key,
                        (unsigned char*)ivecTemp.data(),
                        AES_DECRYPT);

        out = DoPad(out, enc, padding);
        return true;
    }
}















